#!/usr/bin/env python3
# -*- coding: utf-8 -*-
"""
HackerPermKeeper T00ls - Linux环境自动化测试器
功能: 全面测试一键化功能在Linux环境中的表现
作者: HackerPermKeeper Team
"""

import os
import sys
import json
import time
import socket
import subprocess
import threading
from datetime import datetime
from pathlib import Path
import argparse

# 添加项目根目录到Python路径
project_root = Path(__file__).parent.parent
sys.path.insert(0, str(project_root))

try:
    import requests
    import paramiko
    from colorama import init, Fore, Style
    init()  # 初始化colorama
except ImportError as e:
    print(f"❌ 缺少依赖库: {e}")
    print("请安装: pip3 install requests paramiko colorama")
    sys.exit(1)

class Colors:
    """颜色定义"""
    RED = Fore.RED
    GREEN = Fore.GREEN
    YELLOW = Fore.YELLOW
    BLUE = Fore.BLUE
    CYAN = Fore.CYAN
    MAGENTA = Fore.MAGENTA
    WHITE = Fore.WHITE
    RESET = Style.RESET_ALL
    BOLD = Style.BRIGHT

class AutomationTester:
    """自动化测试器主类"""
    
    def __init__(self, target_ip, attacker_ip=None, config_file=None):
        self.target_ip = target_ip
        self.attacker_ip = attacker_ip or self._get_local_ip()
        self.project_root = project_root
        self.test_results = []
        self.start_time = datetime.now()
        
        # 创建必要目录
        self.logs_dir = self.project_root / "logs"
        self.reports_dir = self.project_root / "reports"
        self.logs_dir.mkdir(exist_ok=True)
        self.reports_dir.mkdir(exist_ok=True)
        
        # 日志文件
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        self.log_file = self.logs_dir / f"automation_test_{timestamp}.log"
        
        # 加载配置
        self.config = self._load_config(config_file)
        
        print(f"{Colors.BLUE}{'='*60}{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.CYAN}HackerPermKeeper T00ls 自动化测试器{Colors.RESET}")
        print(f"{Colors.BLUE}{'='*60}{Colors.RESET}")
        print(f"🎯 目标IP: {Colors.YELLOW}{self.target_ip}{Colors.RESET}")
        print(f"🏠 攻击机IP: {Colors.YELLOW}{self.attacker_ip}{Colors.RESET}")
        print(f"📁 项目根目录: {Colors.YELLOW}{self.project_root}{Colors.RESET}")
        print(f"📝 日志文件: {Colors.YELLOW}{self.log_file}{Colors.RESET}")
        print()
    
    def _get_local_ip(self):
        """获取本机IP地址"""
        try:
            s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
            s.connect(("8.8.8.8", 80))
            ip = s.getsockname()[0]
            s.close()
            return ip
        except:
            return "127.0.0.1"
    
    def _load_config(self, config_file):
        """加载配置文件"""
        default_config = {
            "timeouts": {
                "network": 5,
                "ssh": 10,
                "web": 5
            },
            "test_credentials": {
                "ssh": [
                    {"username": "root", "password": "root"},
                    {"username": "root", "password": "toor"},
                    {"username": "root", "password": "123456"},
                    {"username": "admin", "password": "admin"},
                    {"username": "test", "password": "test"}
                ]
            },
            "scan_ports": [21, 22, 23, 25, 53, 80, 110, 443, 993, 995, 3389, 5432, 3306],
            "web_paths": ["/", "/admin", "/login", "/test", "/phpmyadmin", "/wp-admin", "/manager"]
        }
        
        if config_file and Path(config_file).exists():
            try:
                with open(config_file, 'r', encoding='utf-8') as f:
                    user_config = json.load(f)
                    default_config.update(user_config)
            except Exception as e:
                self.log(f"配置文件加载失败: {e}", "warning")
        
        return default_config
    
    def log(self, message, level="info"):
        """日志记录"""
        timestamp = datetime.now().strftime("%Y-%m-%d %H:%M:%S")
        
        # 颜色映射
        color_map = {
            "info": Colors.BLUE,
            "success": Colors.GREEN,
            "warning": Colors.YELLOW,
            "error": Colors.RED,
            "debug": Colors.CYAN
        }
        
        # 符号映射
        symbol_map = {
            "info": "ℹ️",
            "success": "✅",
            "warning": "⚠️",
            "error": "❌",
            "debug": "🔍"
        }
        
        color = color_map.get(level, Colors.WHITE)
        symbol = symbol_map.get(level, "📝")
        
        formatted_message = f"{color}[{timestamp}] {symbol} {message}{Colors.RESET}"
        print(formatted_message)
        
        # 写入日志文件
        with open(self.log_file, 'a', encoding='utf-8') as f:
            f.write(f"[{timestamp}] [{level.upper()}] {message}\n")
    
    def test_network_connectivity(self):
        """测试网络连通性"""
        self.log("测试网络连通性...")
        
        try:
            # 使用socket测试连接
            sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
            sock.settimeout(self.config["timeouts"]["network"])
            result = sock.connect_ex((self.target_ip, 22))  # 测试SSH端口
            sock.close()
            
            if result == 0:
                self.log(f"目标 {self.target_ip} 网络可达", "success")
                self._record_result("网络连通性测试", True, "SSH端口可达")
                return True
            else:
                # 尝试ping
                ping_result = subprocess.run(
                    ["ping", "-c", "3", self.target_ip],
                    capture_output=True, text=True, timeout=10
                )
                if ping_result.returncode == 0:
                    self.log(f"目标 {self.target_ip} 可ping通", "success")
                    self._record_result("网络连通性测试", True, "ping可达")
                    return True
                else:
                    self.log(f"目标 {self.target_ip} 网络不可达", "error")
                    self._record_result("网络连通性测试", False, "网络不可达")
                    return False
        except Exception as e:
            self.log(f"网络测试异常: {e}", "error")
            self._record_result("网络连通性测试", False, str(e))
            return False
    
    def test_port_scan(self):
        """端口扫描测试"""
        self.log("执行端口扫描...")
        
        open_ports = []
        
        for port in self.config["scan_ports"]:
            try:
                sock = socket.socket(socket.AF_INET, socket.SOCK_STREAM)
                sock.settimeout(2)
                result = sock.connect_ex((self.target_ip, port))
                sock.close()
                
                if result == 0:
                    open_ports.append(port)
                    self.log(f"发现开放端口: {port}", "success")
            except:
                continue
        
        if open_ports:
            self.log(f"扫描完成，发现 {len(open_ports)} 个开放端口", "success")
            self._record_result("端口扫描", True, f"开放端口: {open_ports}")
            return open_ports
        else:
            self.log("未发现开放端口", "warning")
            self._record_result("端口扫描", False, "无开放端口")
            return []
    
    def test_ssh_bruteforce(self):
        """SSH暴力破解测试"""
        self.log("测试SSH暴力破解...")
        
        for cred in self.config["test_credentials"]["ssh"]:
            try:
                ssh = paramiko.SSHClient()
                ssh.set_missing_host_key_policy(paramiko.AutoAddPolicy())
                
                ssh.connect(
                    self.target_ip,
                    username=cred["username"],
                    password=cred["password"],
                    timeout=self.config["timeouts"]["ssh"]
                )
                
                # 执行简单命令验证
                stdin, stdout, stderr = ssh.exec_command("whoami")
                result = stdout.read().decode().strip()
                
                ssh.close()
                
                self.log(f"SSH登录成功: {cred['username']}:{cred['password']}", "success")
                self._record_result("SSH暴力破解", True, f"成功凭据: {cred}")
                return True
                
            except paramiko.AuthenticationException:
                continue
            except Exception as e:
                self.log(f"SSH连接异常: {e}", "debug")
                continue
        
        self.log("SSH暴力破解未成功", "warning")
        self._record_result("SSH暴力破解", False, "未发现弱密码")
        return False
    
    def test_web_scan(self):
        """Web服务扫描测试"""
        self.log("测试Web服务扫描...")
        
        web_services = []
        
        for port in [80, 443, 8080, 8443]:
            for protocol in ["http", "https"]:
                if port in [443, 8443] and protocol == "http":
                    continue
                if port in [80, 8080] and protocol == "https":
                    continue
                
                url = f"{protocol}://{self.target_ip}:{port}"
                
                try:
                    response = requests.get(
                        url,
                        timeout=self.config["timeouts"]["web"],
                        verify=False
                    )
                    
                    if response.status_code == 200:
                        web_services.append(url)
                        self.log(f"发现Web服务: {url}", "success")
                        
                        # 测试常见路径
                        self._test_web_paths(url)
                        
                except requests.exceptions.RequestException:
                    continue
        
        if web_services:
            self._record_result("Web服务扫描", True, f"发现服务: {web_services}")
            return web_services
        else:
            self._record_result("Web服务扫描", False, "未发现Web服务")
            return []
    
    def _test_web_paths(self, base_url):
        """测试Web路径"""
        for path in self.config["web_paths"]:
            try:
                url = f"{base_url}{path}"
                response = requests.get(url, timeout=3, verify=False)
                
                if response.status_code == 200:
                    self.log(f"发现可访问路径: {path}", "success")
                elif response.status_code in [301, 302]:
                    self.log(f"发现重定向路径: {path}", "info")
            except:
                continue
    
    def test_auto_decision_system(self):
        """测试自动决策系统"""
        self.log("测试自动决策系统...")
        
        try:
            # 检查自动决策主程序
            auto_decision_path = self.project_root / "apps" / "auto_decision_main.py"
            
            if not auto_decision_path.exists():
                self.log("未找到自动决策主程序", "warning")
                self._record_result("自动决策系统", False, "程序文件不存在")
                return False
            
            # 执行自动决策
            cmd = [
                "python3",
                str(auto_decision_path),
                "--target", self.target_ip,
                "--mode", "scan",
                "--quick"
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=60,
                cwd=str(self.project_root)
            )
            
            if result.returncode == 0:
                self.log("自动决策系统执行成功", "success")
                self._record_result("自动决策系统", True, "扫描模式执行成功")
                return True
            else:
                self.log(f"自动决策系统执行失败: {result.stderr}", "error")
                self._record_result("自动决策系统", False, result.stderr)
                return False
                
        except subprocess.TimeoutExpired:
            self.log("自动决策系统执行超时", "warning")
            self._record_result("自动决策系统", False, "执行超时")
            return False
        except Exception as e:
            self.log(f"自动决策系统测试异常: {e}", "error")
            self._record_result("自动决策系统", False, str(e))
            return False
    
    def test_backdoor_detection(self):
        """测试后门检测功能"""
        self.log("测试后门检测功能...")
        
        try:
            backdoor_check_path = self.project_root / "payload" / "13check" / "backdoor_checks.py"
            
            if not backdoor_check_path.exists():
                self.log("未找到后门检测程序", "warning")
                self._record_result("后门检测", False, "程序文件不存在")
                return False
            
            # 执行后门检测
            cmd = [
                "python3",
                str(backdoor_check_path),
                "--target", self.target_ip,
                "--quick-check"
            ]
            
            result = subprocess.run(
                cmd,
                capture_output=True,
                text=True,
                timeout=30,
                cwd=str(self.project_root)
            )
            
            if result.returncode == 0:
                self.log("后门检测功能正常", "success")
                self._record_result("后门检测", True, "检测功能正常")
                return True
            else:
                self.log(f"后门检测执行失败: {result.stderr}", "warning")
                self._record_result("后门检测", False, result.stderr)
                return False
                
        except Exception as e:
            self.log(f"后门检测测试异常: {e}", "error")
            self._record_result("后门检测", False, str(e))
            return False
    
    def _record_result(self, test_name, success, details):
        """记录测试结果"""
        result = {
            "test": test_name,
            "success": success,
            "details": details,
            "timestamp": datetime.now().isoformat()
        }
        self.test_results.append(result)
    
    def run_all_tests(self):
        """运行所有测试"""
        self.log("开始执行完整测试套件", "info")
        
        tests = [
            ("网络连通性", self.test_network_connectivity),
            ("端口扫描", self.test_port_scan),
            ("SSH暴力破解", self.test_ssh_bruteforce),
            ("Web服务扫描", self.test_web_scan),
            ("自动决策系统", self.test_auto_decision_system),
            ("后门检测", self.test_backdoor_detection)
        ]
        
        for test_name, test_func in tests:
            self.log(f"开始测试: {test_name}")
            try:
                test_func()
            except Exception as e:
                self.log(f"测试 {test_name} 发生异常: {e}", "error")
                self._record_result(test_name, False, f"异常: {e}")
            
            time.sleep(1)  # 测试间隔
        
        self._generate_report()
    
    def _generate_report(self):
        """生成测试报告"""
        self.log("生成测试报告...")
        
        end_time = datetime.now()
        duration = end_time - self.start_time
        
        passed = sum(1 for r in self.test_results if r["success"])
        total = len(self.test_results)
        success_rate = (passed / total * 100) if total > 0 else 0
        
        # JSON报告
        report_data = {
            "test_info": {
                "target_ip": self.target_ip,
                "attacker_ip": self.attacker_ip,
                "start_time": self.start_time.isoformat(),
                "end_time": end_time.isoformat(),
                "duration": str(duration)
            },
            "summary": {
                "total_tests": total,
                "passed": passed,
                "failed": total - passed,
                "success_rate": f"{success_rate:.1f}%"
            },
            "results": self.test_results
        }
        
        # 保存JSON报告
        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        json_report = self.reports_dir / f"automation_test_{timestamp}.json"
        
        with open(json_report, 'w', encoding='utf-8') as f:
            json.dump(report_data, f, indent=2, ensure_ascii=False)
        
        # 生成Markdown报告
        md_report = self.reports_dir / f"automation_test_{timestamp}.md"
        self._generate_markdown_report(md_report, report_data)
        
        # 打印摘要
        print(f"\n{Colors.BLUE}{'='*60}{Colors.RESET}")
        print(f"{Colors.BOLD}{Colors.CYAN}测试完成摘要{Colors.RESET}")
        print(f"{Colors.BLUE}{'='*60}{Colors.RESET}")
        print(f"🎯 目标IP: {Colors.YELLOW}{self.target_ip}{Colors.RESET}")
        print(f"⏱️  测试时长: {Colors.YELLOW}{duration}{Colors.RESET}")
        print(f"📊 总测试数: {Colors.YELLOW}{total}{Colors.RESET}")
        print(f"✅ 通过: {Colors.GREEN}{passed}{Colors.RESET}")
        print(f"❌ 失败: {Colors.RED}{total - passed}{Colors.RESET}")
        print(f"📈 成功率: {Colors.YELLOW}{success_rate:.1f}%{Colors.RESET}")
        print(f"📝 JSON报告: {Colors.CYAN}{json_report}{Colors.RESET}")
        print(f"📄 MD报告: {Colors.CYAN}{md_report}{Colors.RESET}")
        print(f"📋 详细日志: {Colors.CYAN}{self.log_file}{Colors.RESET}")
        print(f"{Colors.BLUE}{'='*60}{Colors.RESET}")
    
    def _generate_markdown_report(self, report_path, data):
        """生成Markdown格式报告"""
        with open(report_path, 'w', encoding='utf-8') as f:
            f.write(f"""# HackerPermKeeper T00ls 自动化测试报告

## 测试信息
- **目标IP**: {data['test_info']['target_ip']}
- **攻击机IP**: {data['test_info']['attacker_ip']}
- **开始时间**: {data['test_info']['start_time']}
- **结束时间**: {data['test_info']['end_time']}
- **测试时长**: {data['test_info']['duration']}

## 测试摘要
- **总测试数**: {data['summary']['total_tests']}
- **通过数**: {data['summary']['passed']}
- **失败数**: {data['summary']['failed']}
- **成功率**: {data['summary']['success_rate']}

## 详细结果

| 测试项目 | 状态 | 详情 | 时间 |
|---------|------|------|------|
""")
            
            for result in data['results']:
                status = "✅ 通过" if result['success'] else "❌ 失败"
                f.write(f"| {result['test']} | {status} | {result['details']} | {result['timestamp']} |\n")
            
            f.write(f"""
## 建议

### 安全建议
1. 如发现SSH弱密码，请立即修改为强密码
2. 如发现开放的Web服务，建议进行安全加固
3. 定期运行安全扫描，及时发现安全隐患

### 测试建议
1. 在生产环境中谨慎使用暴力破解功能
2. 建议在隔离的测试环境中进行完整测试
3. 定期更新测试用例和检测规则

---
*报告由 HackerPermKeeper T00ls 自动生成于 {datetime.now().strftime('%Y-%m-%d %H:%M:%S')}*
""")

def main():
    """主函数"""
    parser = argparse.ArgumentParser(description="HackerPermKeeper T00ls Linux环境自动化测试器")
    parser.add_argument("target_ip", help="目标IP地址")
    parser.add_argument("--attacker-ip", help="攻击机IP地址（可选）")
    parser.add_argument("--config", help="配置文件路径（可选）")
    parser.add_argument("--quick", action="store_true", help="快速测试模式")
    
    args = parser.parse_args()
    
    try:
        tester = AutomationTester(
            target_ip=args.target_ip,
            attacker_ip=args.attacker_ip,
            config_file=args.config
        )
        
        if args.quick:
            # 快速测试模式
            tester.test_network_connectivity()
            tester.test_port_scan()
            tester._generate_report()
        else:
            # 完整测试模式
            tester.run_all_tests()
            
    except KeyboardInterrupt:
        print(f"\n{Colors.YELLOW}测试被用户中断{Colors.RESET}")
    except Exception as e:
        print(f"{Colors.RED}测试过程中发生错误: {e}{Colors.RESET}")

if __name__ == "__main__":
    main()